# coding=utf8

import numpy as np
import pandas as pd
import faker

data_num = 30
df = pd.DataFrame({
        'age': np.random.randint(10, 20, size=data_num),
        'score': np.random.randint(30, 90, size=data_num),
        })

class Preliminary:

    @staticmethod
    def demo_apply():
        df = pd.DataFrame(np.arange(9).reshape((3, 3)),
                          index=list('abc'),
                          # columns=['col1', 'col2', 'col3'],
                          dtype=np.float16)

        def userfun1(ary, axis=0):
            return ary.sum(axis=0)/ary.shape[axis]

        userfun2 = lambda x: np.sin(x)+np.cos(x)

        print(
            "def userfun1(ary, axis=0):\n"
            "    return ary.sum(axis=0)/ary.shape[axis]\n\n"
            "userfun2 = lambda x: np.sin(x)+np.cos(x)\n"
        )

        print(">>> df")
        print(df)

        # 聚合计算
        print(">>> df.apply(userfun1, axis=1)")
        print(df.apply(userfun1, axis=1))

        # 针对DataFrame的逐个元素进行计算
        print(">>> df.apply(userfun2)")
        print(df.apply(userfun2))

        # apply用于Series时，为针对逐个元素的计算
        # 使用userfun1时，触发异常AttributeError: 'float' object has no attribute 'sum'
        print(">>> df[0].apply(userfun2)")
        print(df[0].apply(userfun2))

    @staticmethod
    def demo_apply_result_type():
        """
        result_type: 字符串, 或None。取值：{'expand', 'reduce', 'broadcast', None}. 缺省为None。
        用于确定函数返回值为列表或Series时处理的方式。
        该选项参数只对axis=1时起作用。
        各个选项的计算方式：
        * 'expand': 函数值为列表时，展开为列数据。
        * 'reduce': 函数值为Series时，不展开为列数据。
        * 'broadcast': 函数值以广播方式沿所指定的轴向传递给原数据集，保持原有索引或列标签。
          该方式要求结果长度与相应轴向数据长度相同。
        * None：处理方式依赖于函数返回值，返回值为列表时以Series结构处理，不展开为列数据；
          返回值为Series时，展开为列数据。
        """

        # 通用函数
        def f1(x):
            return [np.exp(x), np.log(x)]

        # 聚合函数
        def f2(x):
            return [np.mean(x), np.sum(x)]

        # 聚合函数，返回类型为Series
        def f3(x):
            return pd.Series([np.mean(x), np.var(x)], index=['mean', 'var'])

        print("f1->list, f2->list, f3->Series")
        print(f">>> df\n{df}")

        # result_type='expand'
        # 该方式下， 函数返回值为序列时，将展开为列
        print("---result_type = expand---")
        print("# 该方式下，函数返回值为序列时，将展开为列")
        print("# 1）通用函数：列表展开为列，列名为整数索引")
        print(">>> f1 = lambda x: [np.exp(x), np.log(x)]")
        print(">>> df.apply(func=f1, axis=1, result_type='expand')")
        df_expand_list = df.apply(func=f1, axis=1, result_type='expand')
        print(df_expand_list.head(2))
        print("# 2）聚合函数：列表展开为列，列名为整数索引")
        print(">>> f2 = lambda x: [np.mean(x), np.var(x)]")
        print(">>> df.apply(func=f1, axis=1, result_type='expand')")
        df_expand_list2 = df.apply(func=f2, axis=1, result_type='expand')
        print(df_expand_list2.head(2))
        print("# 3）聚合函数：Series展开为列，列名为Series索引")
        print(">>> f3 = lambda x: pd.Series([np.mean(x), np.var(x)], index=['mean', 'var'])")
        print(">>> df.apply(func=f3, axis=3, result_type='expand')")
        df_expand_series = df.apply(func=f3, axis=1, result_type='expand')
        print(df_expand_series.head(2))

        # result_type='reduce'
        print("---result_type = reduce---")
        print("# 非聚合函数：")
        print(">>> f1 = lambda x: [np.exp(x), np.log(x)]")
        print(">>> df.apply(func=f1, axis=1, result_type='reduce')")
        df_reduce = df.apply(func=f1, axis=1, result_type='reduce')
        print(df_reduce.head(2))
        print("# 聚合函数：")
        print(">>> f2 = lambda x: [np.mean(x), np.sum(x)]")
        print(">>> df.apply(func=f2, axis=1, result_type='reduce')")
        df_reduce = df.apply(func=f2, axis=1, result_type='reduce')
        print(df_reduce.head(2))

        # result_type='broadcast'
        # 函数值以广播方式沿所指定的轴向传递给原数据集，保持原有索引或列标签。
        # 1）
        # columns length must be as same as result length of user_fun
        # 该方式要求：函数返回值长度与列标签长度相同。
        # otherwise, ValueError will be raised!
        # 否则就会触发异常：ValueError: cannot broadcast result
        # df.loc[:, 'all'] = df.age + df.score
        # print(df.head(3))
        # 2）
        # 如果返回序列与列标签长度相同，替换其值，标签名称不变
        print("--- result_type = broadcast ---")
        print("# 如果返回序列与列标签长度相同，替换其值，标签名称不变")
        print(">>> f2 = lambda x: [np.mean(x), np.sum(x)])")
        print(">>> df.apply(func=f2, axis=1, result_type='broadcast')")
        df_broadcast2 = df.apply(func=f2, axis=1, result_type='broadcast')
        print(df_broadcast2.head(2))
        print(">>> f3 = lambda x: pd.Series([np.mean(x), np.var(x)], index=['mean', 'var'])")
        print(">>> df.apply(func=f3, axis=1, result_type='broadcast')")
        df_broadcast3 = df.apply(func=f3, axis=1, result_type='broadcast')
        print(df_broadcast3.head(2))

        # result_type='None'
        # 1）
        # 函数返回值为列表，apply结果为Series
        # apply return Series if user_fun return list
        df_none_list = df.apply(func=f2, axis=1, result_type=None)
        print("--- result_type = None ---")
        print("# set to Series if user_fun return list")
        print(">>> df.apply(func=f2, axis=1, result_type=None)")
        print(df_none_list.head(2))
        # 2）
        # 函数返回值为Series，apply结果为DataFrame
        # apply return DataFrame if user_fun return Series
        df_none_series = df.apply(func=f3, axis=1, result_type=None)
        print("---df_none_Series---")
        print("# expand to columns if user_fun return Series")
        print(df_none_series.head(2))

        # 在axis=0方向，不允许使用
        try:
            print("# 使用axis=0, 会触发ValueError")
            print(">>> df.apply(func=f3, axis=0, result_type='broadcast')")
            df_broadcast2 = df.apply(func=f3, axis=0, result_type='broadcast')
            print(df_broadcast2)
        except Exception as e:
            print(e)

    @staticmethod
    def demo_agg():
        df = pd.DataFrame(np.arange(9).reshape((3, 3)), index=list('abc'), dtype=np.float16)

        print(">>> df")
        print(df)

        # 调用单个函数
        # 1)聚合函数
        print(">>> fun11 = lambda x: np.max(x)-min(x)")
        print(">>> df1 = df.agg(fun11, axis=1)")
        print(">>> df1")
        fun11 = lambda x: np.max(x)-min(x)
        df11 = df.agg(fun11, axis=1)
        print(df11)
        # 2）一般函数
        print(">>> fun12 = lambda x: np.exp(x)")
        print(">>> df12 = df.agg(fun12, axis=1)")
        print(">>> df12")
        fun12 = lambda x: np.exp(x)
        df12 = df.agg(fun12, axis=1)
        print(df12)

        # 调用多个函数
        # 1）聚合函数（内置聚合函数）
        print(">>> fun2 = lambda x: np.max(x)/np.mean(x)")
        print(">>> df21 = df.apply([fun11, fun2], 1, 10)")
        print(">>> df21")
        df21 = df.agg([np.max, np.median], 1)
        print(df21)
        # 2) 一般函数
        df22 = df.agg([fun11])
        print(">>> df22 = df.agg([fun11])")
        print(df22)
        # 3） 不支持一般处理函数函数组合
        # fun22 = lambda z: np.var(z)
        # df23 = df.agg([fun11, fun22])
        # ValueError: cannot combine transform and aggregation operations

        # 3）传递参数
        print(">>> fun3 = lambda x, y=9: np.max(x)+y")
        print(">>> df3 = df.apply(fun3, axis=1, y=20)")
        print(">>> df3")
        fun3 = lambda x, y=1: np.max(x)+y
        df3 = df.agg(fun3, axis=1, y=20)

        # 4）使用字典
        # 指定行调用函数（axis=0）
        print(">>> df.agg({'a': np.max, 'b': np.min}, axis=1)")
        df41 = df.agg({'a': np.max, 'b': np.min}, axis=1)
        print(df41)
        # 指定列调用函数（axis=1）
        print(">>> df.agg({'a': np.max, 'b': 'sum', 'c': 'median'}, axis=1)")
        df42 = df.agg({'a': np.max, 'b': "sum", 'c': 'median'}, axis=1)
        print(df42)

    @staticmethod
    def demo_agg_result():
        # Series， 调用单个聚合函数, 结果为标量
        sr = pd.Series(range(5))
        print(">>> sr")
        print(sr)
        print(">>> sr.agg(np.max)")
        print(sr.agg(np.max))

        # Series， 调用单个通用函数, 结果为Series
        print(">>> sr.agg(np.sin)")
        print(sr.agg(np.sin))

        # Series， 调用多个聚合函数, 结果为Series
        print(">>> sr.agg([np.max, np.min])")
        print(sr.agg([np.max, np.min]))

        # Series， 调用多个通用函数, 结果为DataFrame
        print(">>> sr.agg([np.sin, np.cos])")
        print(sr.agg([np.sin, np.cos]))

        # DataFrame， 调用单个聚合函数, 结果为Series
        df = pd.DataFrame(np.arange(9).reshape((3, 3)), columns=list('abc'))
        print(">>> df")
        print(df)
        print(">>> df.agg(np.max)")
        print(df.agg(np.max))

        # DataFrame， 调用多个聚合函数, 结果为DataFrame
        print(">>> df.agg([np.max, np.mean, np.min])")
        print(df.agg([np.max, np.mean, np.min]))

        # DataFrame， 调用多个通用函数, 每个数据产生多个结果，结果为DataFrame，在另一轴向增加一层索引
        print(">>> df.agg([np.sin, np.cos])")
        print(df.agg([np.sin, np.cos]))
        print(">>> df.agg([np.sin, np.cos], axis=1)")
        print(df.agg([np.sin, np.cos], axis=1))


class Task:

    @classmethod
    def user_fun(cls, data=df):

        def fun_age_bool(x):
            if 12 < x < 14 or x < 11 or x > 16:
                return True
            else:
                return False
        age_bool = data.age.apply(fun_age_bool)
        # age_bool = ((data.age > 12) & (data.age < 14)) | (data.age < 11) | (data.age > 16)
        data2 = data.loc[age_bool]

        def f(x):
            if x < 17:
                return 15
            elif x == 17 or x == 18:
                return 10
            else:
                return 5

        # data3 = data2.score + data2.age.apply_result_type(f)
        data3 = data2.score + data2.agg({'age': f})['age']
        special_mean = data3.mean()
        print('自定义函数统计结果：{:.2f}'.format(special_mean))
        Task.test_user_fun()

    @classmethod
    def test_user_fun(cls, data=df):

        def fun_age_bool(x):
            if 12 < x < 14 or x < 11 or x > 16:
                return True
            else:
                return False

        age_bool1 = data.age.apply(fun_age_bool)
        age_bool2 = ((data.age > 12) & (data.age < 14)) | (data.age < 11) | (data.age > 16)
        print("test bool1 == bool2:", np.all(age_bool1 == age_bool2))
        data2 = data.loc[age_bool1]
        print(data2.head(3))

        # use bool value as 0, 1
        df15 = (data2.age < 17)*15
        df10 = ((data2.age == 17) | (data2.age == 18))*10
        df5 = (data2.age > 18)*5
        data31 = data2.score + df15 + df10 + df5
        sp_mean1 = data31.mean()
        print("add score by logic expression:")
        print(data31.head(3))

        def f(x):
            if x < 17:
                return 15
            elif x == 17 or x == 18:
                return 10
            else:
                return 5

        # data32 = data2.score + data2.agg(func={'age': f}).age
        data32 = data2.score + data2.age.apply(f)
        special_mean = data32.mean()
        print(
            "add score result is same by logic expression and user function:"
            f"{np.all(data31 == data32)}"
        )

        print('自定义函数统计结果2：{:.2f}'.format(sp_mean1))
        print('自定义函数统计结果3：{:.2f}'.format(special_mean))


if __name__ == "__main__":
    Preliminary.demo_apply()
    # Preliminary.demo_agg()
    # Preliminary.demo_agg_result()
    # Preliminary.demo_apply_result_type()
    # Task.user_fun(df)
    # exercise_1()
    # exercise_2()
